#pragma once

#include "core/packet.h"
#include "utils/spsc_queue.h"

#include <atomic>
#include <memory>
#include <vector>
#include <functional>

/**
 * @brief Packet distribution strategy
 */
enum class DispatchStrategy {
  ROUND_ROBIN, /// Round-robin distribution has the lowest overhead
  FLOW_HASH, /// Use flow-based hashing to ensure in-order processing within the
             /// same flow
  LOAD_AWARE, /// Distribution based on queue load awareness
  ADAPTIVE    /// Adaptive strategy switching
};

/**
 * @brief Inteligent packet dispatcher
 *
 * An SPSC queue responsible for distributing captured packets to multiple
 * worker threads. Supports multiple dstribution strategies and load balancing.
 */
class PacketDispatcher {
 public:
    using PacketQueue = utils::SPSCQueue<Packet, 4096>;
    using QueueStatsCallback = std::function<void(size_t queue_index, size_t queue_size)>;

    /**
     * @brief Constructor
     * @param queue_count The number of queue (worker thread)
     * @param strategy Initial distribution strategy
     */
    PacketDispatcher(size_t queue_count, DispatchStrategy strategy = DispatchStrategy::ROUND_ROBIN);
    ~PacketDispatcher();

    /// Disable copying and assignment
    PacketDispatcher(const PacketDispatcher&) = delete;
    PacketDispatcher& operator=(const PacketDispatcher&) = delete;

    /**
     * @brief Distribute packets to the appropriate queue
     * @param packet Packets to be distributed (move)
     * @return Returns true on successful distribution, false if all queues are full
     */
    bool dispatch(Packet&& packet);

    /**
     * @brief Batch packet distribution
     * @param packets Array of packets
     * @param count Number of packets
     * @return Number of packets successfully distributed
     */
    size_t dispatchBulk(Packet* packets, size_t count);

    void setStrategy(DispatchStrategy strategy);

    DispatchStrategy getStrategy() const { return m_strategy.load(std::memory_order::acquire); }

    /**
     * @brief Get the queue at the specified index
     */
    PacketQueue& getQueue(size_t index);

    size_t getQueueCount() const { return m_queues.size(); }

    /**
     * @brief Set a queue status callback for monitoring
     */
    void setStatsCallback(QueueStatsCallback callback) { m_statsCallback = callback; }

    /**
     * @brief Statistics
     */
    struct DispatchStats {
        uint64_t total_dispatched;
        uint64_t total_failures;
        uint64_t strategy_changes;
        std::vector<uint64_t> queue_dispatched;
        std::vector<size_t> queue_sizes;
        double average_decision_time_ns;
    };

    DispatchStats getStats() const;

    void resetStats();

 private:
    /*
    /// Inner queue structure
    struct QueueInfo {
        std::unique_ptr<PacketQueue> queue;
        std::atomic<size_t> dispatched_count;
        std::atomic<size_t> current_size;

        QueueInfo() : queue(std::make_unique<PacketQueue>()), dispatched_count(0), current_size(0) {}

        QueueInfo(const QueueInfo&) = delete;
        QueueInfo(QueueInfo&&) = delete;
        QueueInfo& operator=(const QueueInfo&) = delete;
        QueueInfo& operator=(QueueInfo&&) = delete;
    };
    */

    /// Distribution strategy implementation
    size_t selectQueueRoundRobin();
    size_t selectQueueFlowHash(const Packet& packet);
    size_t selectQueueLoadAware();
    size_t selectQueueAdaptive(const Packet& packet);

    uint32_t calculateFlowHash(const Packet& packet) const;
    size_t findLeastLoadedQueue() const;
    // void updateQueueStats(size_t queue_index);

    std::vector<std::shared_ptr<PacketQueue>> m_queues;
    std::atomic<DispatchStrategy> m_strategy;
    std::atomic<size_t> m_roundRobinIndex{0};

    mutable std::mutex m_statsMutex;
    DispatchStats m_stats;
    std::atomic<uint64_t> m_totalDecisionTime{0};

    QueueStatsCallback m_statsCallback;
};
